# -*- coding: utf-8 -*-
"""
Created on xxx  2024-2025
PyTorch编程实现循环神经网络的英文文本生成
参考代码： P316-324页
数据集：安徒生童话（英文：皇帝的新装）
实验环境：python3.9.13、torch1.10.2+cu113
数据处理：字符变成索引，再通过embedding层变成张量
神经网络：embedding+lstm+linear
实验目的：英文字符转数字
实验要求：loss，150字生成短文
代码特点：温度越高--抽样分布熵更大，更容易生成出人意料的句子，温度越低--随机性更小，生成结果可预测性增大
        较低温度（如 0.2）：输出更确定（更常见）
        较高温度（如 1.2）：输出更随机（更多样性）
@author: htq
"""
# -------------------- 导入包库 --------------------
import random
import torch
import torch.nn as nn
import string
import gzip
import time
import matplotlib.pyplot as plt
import math

import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'  # 允许重复加载动态链接库


# -------------------- 参数定义 --------------------
batch_size = 64    # 每次训练使用的样本数
n_iters = 3000     # 总的训练迭代次数
print_every = 500  # 每隔多少轮打印一次训练信息
plot_every = 10    # 每隔多少轮绘制一次损失曲线
hidden_size = 100  # LSTM隐藏层的大小
num_layers = 1     # LSTM的层数
temperature_list = [0.2, 0.5, 1.0, 1.2]  # 控制生成文本时的温度参数列表
generated_len = 200   # 生成文本的长度


# -------------------- 数据处理 --------------------
def random_chunk(corpus, corpus_len, chunk_len):
    """随机截取语料库并返回固定长度的文本
    corpus：一个字符串，表示完整的语料库
    corpus_len：一个整数，表示语料库长度
    chunk_len：一个整数，表示希望截取的文本块长度
    """
    start_idx = random.randint(0, corpus_len - chunk_len)  # 确保截取的文本块不超过语料库的长度
    end_idx = start_idx + chunk_len + 1  # 确保截取的文本块包含chunk_len个字符
    return corpus[start_idx: end_idx]


def char_tensor(my_str, all_chars):
    """将字符串转为索引变量
    my_str：需要转换的字符串
    all_chars：包含所有可能的字符列表或字符串，用于构建索引
    """
    tensor_idx = torch.zeros((len(my_str),), dtype=torch.long)  # 创建一个与输入字符长度相同的零向量，数据类型为长整型
    for c in range(len(my_str)):
        tensor_idx[c] = all_chars.index(my_str[c])
    return tensor_idx


def make_training_sample(corpus, corpus_len, chunk_len, all_chars):
    """构建一个训练样本，从给定的语料库中随机截取一段文本，并将其转为适合训练的输入和目标张量
    Mnay years ago, there was an Emperor.
    data: Many years ago, there was an Emperor
    target: any years ago, there was an Emperor.
    """
    chunk = random_chunk(corpus, corpus_len, chunk_len)
    data = char_tensor(chunk[:-1], all_chars)  # 从chunk中提取前chunk_len-1个字符（即从第一个字符到倒数第二个字符）的索引张量
    target = char_tensor(chunk[1:], all_chars)  # 从chunk中提取后chunk_len-1个字符（即从第二个字符到最后一个字符）的索引张量
    return data, target


def time_since(since):
    s = time.time() - since
    m = math.floor(s / 60)
    s -= m * 60
    return f"{m:2.0f}分{s:2.0f}秒"


# -------------------- 网络搭建 --------------------
class LSTM(nn.Module):
    def __init__(self, input_size, hidden_dim, output_size, num_layers=1):
        super(LSTM, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_dim
        self.output_size = output_size
        self.num_layers = num_layers
        self.emb = nn.Embedding(input_size, hidden_dim)  # 嵌入层，将输入的整数索引转为固定大小的向量
        self.lstm = nn.LSTM(hidden_dim, hidden_dim, num_layers)  # LSTM层，用于处理序列数据
        self.fc = nn.Linear(hidden_dim, output_size)  # 全连接层，将LSTM的输出映射到最终输出的大小，起分类作用

    def forward(self, inp, hidden):
        """
        inp:输入序列
        hidden:隐藏层状态
        """
        seq_len = len(inp)
        out = self.emb(inp.view(seq_len, 1))
        output, hidden = self.lstm(out.view(seq_len, 1, -1), hidden)
        output = self.fc(output.view(seq_len, -1))
        return output, hidden

    def init_hidden(self):
        """
        初始化隐藏层和细胞状态
        """
        h0 = torch.zeros((self.num_layers, 1, self.hidden_size), requires_grad=True)
        c0 = torch.zeros((self.num_layers, 1, self.hidden_size), requires_grad=True)
        return h0, c0


def train(model, criterion, optimizer, inp, target):
    # model.train()  # 开启训练模式
    hidden = model.init_hidden()  # 初始化隐藏状态，确保每次训练时从相同的初始状态开始
    model.zero_grad()
    output, hidden = model(inp, hidden)
    loss = criterion(output, target)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1)  # 梯度裁剪，防止梯度爆炸。将梯度范围限制在1以内，防止梯度过大
    optimizer.step()
    return loss.item()


# -------------------- 文本生成 --------------------
def generate(model, preset_str, all_chars, predict_len=100, temperature=0.8):
    """使用训练好的模型随机生成文本
    model：训练好的模型
    preset_str:初始的字符串，作为文本的前缀 'Th'
    all_chars:所有的字符列表或字符串
    predict_len:要生成的文本长度
    temperature:控制输出分布的平滑程度，值越高越随机，越低越确定
    """
    hidden = model.init_hidden()
    preset_input = char_tensor(preset_str, all_chars)
    predicted = preset_str

    model.eval()  # 开启评估模型
    with torch.no_grad():
        # 遍历原始字符串中的每个字符（除了最后一个），更新隐藏状态
        for idx in range(len(preset_str) - 1):
            _, hidden = model(preset_input[idx].unsqueeze(0), hidden)

        inp = preset_input[-1]  # 将给定字符串的最后一个字符作为预测的第一个输入

        # 迭代predict_len次，每次生成一个新的字符并更新隐藏状态
        for idx in range(predict_len):
            output, hidden = model(inp.unsqueeze(0), hidden)

            # 将模型的输出概率进行温度缩放，然后使用torch.multinomial根据概率分布采样一个字符索引
            # output.data.view(-1) 将输出展平为一维张量
            # .div(temperature) 将每个元素除以温度参数 temperature。较高的温度值会使输出更加随机，较低的温度值会使输出更加确定
            # .exp() 计算每个元素的指数值，从而将 logits 转换为概率分布
            output_dist = output.data.view(-1).div(temperature).exp()
            top_i = torch.multinomial(output_dist, 1)[0]

            # 将采用得到的字符添加到预测字符串中，并将该字符转换为张量形式，作为下一个时间步的输入
            predicted_char = all_chars[top_i]
            predicted += predicted_char
            inp = char_tensor(predicted_char, all_chars)
    return predicted


# -------------------- 主函数 --------------------
def main():
    all_chars = string.printable  # 英文：A-z a-z 数字 标点符号 其它可以打印的字符
    n_chars = len(all_chars)

    # 读取文件,打印语料库长度
    # corpus = gzip.open("../datasets/shakespeare.txt.gz", 'rt', encoding='utf-8').read()
    with open('./trailData/antushengtonghua.txt', 'rt', encoding='utf-8') as file:
        corpus = file.read()
    corpus_len = len(corpus)
    print("成功读取文本文件，语料长度为：", corpus_len)

    # 随机截断文件
    chunk_len = 200
    print("随机截断文件结果:", random_chunk(corpus, corpus_len, chunk_len))

    # 实例化网络模型，损失函数，优化器
    model = LSTM(n_chars, hidden_size, n_chars, num_layers)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    print("\n开始训练...")
    start = time.time()  # 开始计时
    all_losses = []
    loss_accumulated = 0

    for it in range(1, n_iters + 1):
        loss = train(model, criterion, optimizer, *make_training_sample(corpus, corpus_len, chunk_len, all_chars))
        loss_accumulated += loss

        if it % print_every == 0:
            print(f"[第{it:2d}轮/共{n_iters:4d}轮 进度：{it / n_iters:.2%}] 耗时：{time_since(start):s} 损失:{loss:.4f}")

            for temp in temperature_list:
                print("当前温度：", temp)
                print(generate(model, 'Th', all_chars, temperature=temp), '\n')

            start = time.time()  # 重新计时

        if it % plot_every == 0:
            all_losses.append(loss_accumulated / plot_every)
            loss_accumulated = 0

    plt.figure()
    plt.plot(all_losses)
    plt.title("Losses History")
    plt.show()


if __name__ == '__main__':
    main()
